"""
通过selective search 生成候选框并且判断那些是正例那些是负例,然后生成用于classification的数据集

负例的判断标准与论文中相同定为0.3

"""
import shutil
import time
import cv2
import numpy as np

from utils import check_dir, parse_xxx_csv, parse_anno_xml2loc
from pascal_voc_data_Split2TrainVal import  get_classes  # 文件夹检测及获取所有的分类名
import selective_search  # 生成候选框

# 不同分类 train /val的基路径
BASE_PATH = './data/VOC_{}/{}/'
CLASSIFY_PATH = './data/VOC_{}_classifier/{}/'

#  的anno/img路径
BASE_ANNO_PATH = BASE_PATH + 'Annotations/'
BASE_JPEG_PATH = BASE_PATH + 'JPEGImages/'

# classifier train /val 的anno/img路径
CLASSIFY_ANNO_PATH = CLASSIFY_PATH + 'Annotations/'
CLASSIFY_JPEG_PATH = CLASSIFY_PATH + 'JPEGImages/'


# =============================================================
def generate_classifier():

    # 得到所有的name名
    class_names = get_classes()
    # 遍历所有分类名,为对应分类名创建classifier文件夹
    for name in class_names:
        generate_one_class(name)


# =============================================================
def generate_one_class(name):
    # 操作训练集与测试集
    for data_type in ['train', 'val']:
        # 检测目标文件夹是否存在(不存在创建)
        src_base_dir = CLASSIFY_JPEG_PATH.format(name, data_type)
        check_dir(src_base_dir)
        # 拷贝train/val中的csv文件
        shutil.copy(BASE_PATH.format(name, data_type) + name + '.csv',
                    CLASSIFY_JPEG_PATH.format(name, data_type) + name + '.csv')

        # anno/jpeg文件路径
        src_anno_dir = BASE_ANNO_PATH.format(name, data_type)
        trg_anno_dir = CLASSIFY_ANNO_PATH.format(name, data_type)
        check_dir(trg_anno_dir)

        src_jpeg_dir = BASE_JPEG_PATH.format(name, data_type)
        trg_jpeg_dir = CLASSIFY_JPEG_PATH.format(name, data_type)
        check_dir(trg_jpeg_dir)

        # 计数
        total_num_positive = 0
        total_num_negative = 0

        # 获取该分类train/val样本csv文件
        samples = parse_xxx_csv(src_base_dir + name + '.csv')

        # 处理每一张图片
        for sample_name in samples:
            try:
                since = time.time()
                # 将csv中的文件名称填充成6位
                sample_name = sample_name.zfill(6)
                # 取得源标注path
                src_anno_one_path = src_anno_dir + sample_name + '.xml'
                # 获取源图片path
                src_jpeg_one_path = src_jpeg_dir + sample_name + '.jpg'
                trg_jpeg_one_path = trg_jpeg_dir + sample_name + '.jpg'
                # 通过iou计算得到正例和负例的候选框信息
                positive_list, negative_list = generate_iou(
                    src_anno_one_path,
                    src_jpeg_one_path,

                    name
                )
                # 保存正例和负例
                np.savetxt(trg_anno_dir + sample_name + '_1.csv', np.array(positive_list), fmt='%d', delimiter=' ')
                np.savetxt(trg_anno_dir + sample_name + '_0.csv', np.array(negative_list), fmt='%d', delimiter=' ')
                # 保存图片
                shutil.copyfile(src_jpeg_one_path, trg_jpeg_one_path)
                total_num_positive += len(positive_list)
                total_num_negative += len(negative_list)
                print('parse {} {}.jpeg on {:.0f}m:{:.0f}s'.format(name, sample_name, (time.time() - since) // 60,
                                                                   (time.time() - since) % 60))
            except Exception as e:
                print(e.with_traceback())
                continue
# ===============================================================


def generate_one_img():
    pass


# =============================================================
def generate_iou(anno_path, jpeg_path, name):
    """
    通过selective search 生成候选框,然后通过iou算法决定正例(正例来自anno)和负例

    现在的思路是 得到候选框坐标,然后进行求交,若交集面积小于0.3 则判定为负例
        (为了进一步限制负样本数据,负样本候选框面积必须大于原标注框的1/5)

    :param anno_path:
    :param jpeg_path:
    :param gs:
    :param name: 分类名称
    :return: 正例坐标,负例坐标
    """
    positive_list = []
    negative_list = []
    # 读取anno信息
    positive_list = parse_anno_xml2loc(anno_path, name)
    # 根据得到的bonding box获取标注框中的最大面积(用于后续比较负例)
    max_bndbox_size = get_max_size(positive_list)
    # 生成候选框信息
    rect = get_rects(jpeg_path)
    print('rects len is {}'.format(len(rect)))
    #  计算iou
    iou_list = compute_iou(rect, positive_list)
    # 统计iou值在(0-0.3]之间且面积大于max_bndbox_size/5(最大候选框面积五分之一)的矩形区域成为负样本
    for i in range(len(iou_list)):
        # 计算矩形面积
        rect_area = get_area(rect[i])
        if 0 < iou_list[i] <= 0.3 and rect_area >= (max_bndbox_size/5):
            negative_list.append(rect[i])

    return positive_list, negative_list


# =============================================================

def get_max_size(positive_list):
    """
    计算bounding box 中面积最大的
    :param positive_list:
    :return:
    """
    max_bndbox_size = 0
    for item in positive_list:
        for i in range(len(item)):
            item[i] = int(item[i])
        temp = (item[2] - item[0]) * (item[3] - item[1])
        if temp > max_bndbox_size:
            max_bndbox_size = temp
    return max_bndbox_size
# =============================================================


def get_rects(jpeg_path):
    """
    生成候选框
    :param jpeg_path:
    :param gs:
    :return:
    """
    # 初始化gs
    gs = selective_search.get_selective_search()
    img = cv2.imread(jpeg_path, cv2.IMREAD_COLOR)  # 读取图片
    selective_search.config(gs, img, strategy='q')  # 设置生成策略
    rect = selective_search.get_rects(gs)  # 得到候选框
    return rect
# =============================================================


def compute_iou(rects, bounding_boxs):
    """
    计算iou根据传入的矩形以及候选框,计算矩形的iou值
    :param rects:
    :param bounding_boxs:
    :return:
    """
    iou_list = []
    for rect in rects:
        max_iou = 0
        # 循环候选框,若交集为空iou=0,若交集不为空计算其大小得到iou的值 = 相交面积/(两矩形的面积-相交面积)
        for box in bounding_boxs:
            iou = get_iou(rect, box)
            if iou > max_iou:
                max_iou = iou
        iou_list.append(max_iou)
    return iou_list
# =============================================================


def get_iou(rect, box):
    """
    计算矩形框与候选框iou
    :param rect:
    :param box:
    :return:
    """
    ax1, ay1, ax2, ay2 = rect
    bx1, by1, bx2, by2 = box
    if ax2 <= bx1 or bx2 <= ax1:
        return 0
    if ay2 <= by1 or by2 <= ay1:
        return 0
    area = (min(ax2, by2)-max(ax1, bx1)) * (min(ay2, by2)-max(ay1, by1))
    a_area = get_area(rect)
    b_area = get_area(box)
    return area/(a_area+b_area-area)


def get_area(rect):
    return (rect[2]-rect[0]) * (rect[3]-rect[1])
# =============================================================


if __name__ == '__main__':
    #generate_classifier()
    generate_one_class('car')